import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
import os

# 图片路径
image_path = "input_images/普京-图片.jpg"

# 确保输入图片目录存在
if not os.path.exists(image_path):
    os.makedirs(os.path.dirname(image_path), exist_ok=True)
    print(f"警告: 图片 {image_path} 不存在，请确保图片路径正确")

# 创建FaceAnalysis对象
app = FaceAnalysis(providers=['CPUExecutionProvider'])
app.prepare(ctx_id=-1, det_size=(640, 640))

# 读取图片
try:
    image = cv2.imread(image_path)
    assert image is not None, f"无法读取图片，请检查图片路径是否正确: {image_path}"
except Exception as e:
    print(f"读取图片时出错: {e}")
    # 使用insightface提供的示例图片作为备选
    from insightface.data import get_image as ins_get_image
    print("使用insightface提供的示例图片替代...")
    image = ins_get_image('t1')

# 执行人脸检测
faces = app.get(image)
print(f"检测到 {len(faces)} 个人脸")

# 在图片上绘制检测结果
draw = image.copy()
for idx, face in enumerate(faces):
    # 获取人脸边界框坐标
    bbox = face.bbox.astype(int)
    x1, y1, x2, y2 = bbox
    
    # 绘制边界框
    cv2.rectangle(draw, (x1, y1), (x2, y2), (0, 255, 0), 2)
    
    # 绘制关键点
    if face.kps is not None:
        for kp in face.kps:
            x, y = kp
            cv2.circle(draw, (int(x), int(y)), 1, (0, 0, 255), 2)
    
    # 显示性别和年龄信息（如果有）
    if hasattr(face, 'gender_age'):
        gender, age = face.gender_age
        gender_text = '男' if gender == 1 else '女'
        cv2.putText(draw, f"{gender_text}, {age:.1f}岁", (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

# 保存结果
output_dir = "output_images"
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, "result.jpg")
cv2.imwrite(output_path, draw)
print(f"结果已保存到: {output_path}")

# 显示结果（可选，取决于环境）
try:
    cv2.imshow("Face Detection Result", draw)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
except:
    print("无法显示图像，但结果已保存到文件")